import asyncio
import base64
import json
from typing import Generic, Optional, TypeVar

from browser_use import Browser as BrowserUseBrowser
from browser_use import BrowserConfig
from browser_use.browser.context import BrowserContext, BrowserContextConfig
from browser_use.dom.service import DomService
from pydantic import Field, field_validator
from pydantic_core.core_schema import ValidationInfo

from app.config import config
from app.llm import LLM
from app.tool.base import BaseTool, ToolResult
from app.tool.web_search import WebSearch


_BROWSER_DESCRIPTION = """
与网页浏览器交互以执行各种操作,如导航、元素交互、内容提取和标签页管理。此工具提供全面的浏览器自动化功能:

导航:
- 'go_to_url': 在当前标签页中访问指定URL
- 'go_back': 返回上一页
- 'refresh': 刷新当前页面
- 'web_search': 在当前标签页中搜索查询,查询应该像人类在网上搜索一样,具体且不含糊或过长。只包含最重要的项目。

元素交互:
- 'click_element': 通过索引点击元素
- 'input_text': 在表单元素中输入文本
- 'scroll_down'/'scroll_up': 滚动页面(可选指定像素数)
- 'scroll_to_text': 如果找不到要交互的内容,滚动到该内容
- 'send_keys': 发送特殊键字符串,如Escape、Backspace、Insert、PageDown、Delete、Enter,也支持快捷键如`Control+o`、`Control+Shift+T`等。这用于keyboard.press。
- 'get_dropdown_options': 获取下拉列表中的所有选项
- 'select_dropdown_option': 通过选项文本为交互元素索引选择下拉选项

内容提取:
- 'extract_content': 提取页面内容以获取页面中的特定信息,例如所有公司名称、特定描述、所有相关信息,以结构化格式显示的公司链接或简单链接

标签页管理:
- 'switch_tab': 切换到特定标签页
- 'open_tab': 用URL打开新标签页
- 'close_tab': 关闭当前标签页

实用工具:
- 'wait': 等待指定秒数
"""

Context = TypeVar("Context")


class BrowserUseTool(BaseTool, Generic[Context]):
    name: str = "browser_use"
    description: str = _BROWSER_DESCRIPTION
    parameters: dict = {
        "type": "object",
        "properties": {
            "action": {
                "type": "string",
                "enum": [
                    "go_to_url",
                    "click_element",
                    "input_text",
                    "scroll_down",
                    "scroll_up",
                    "scroll_to_text",
                    "send_keys",
                    "get_dropdown_options",
                    "select_dropdown_option",
                    "go_back",
                    "web_search",
                    "wait",
                    "extract_content",
                    "switch_tab",
                    "open_tab",
                    "close_tab",
                ],
                "description": "要执行的浏览器操作",
            },
            "url": {
                "type": "string",
                "description": "'go_to_url'或'open_tab'操作的URL",
            },
            "index": {
                "type": "integer",
                "description": "'click_element'、'input_text'、'get_dropdown_options'或'select_dropdown_option'操作的元素索引",
            },
            "text": {
                "type": "string",
                "description": "'input_text'、'scroll_to_text'或'select_dropdown_option'操作的文本",
            },
            "scroll_amount": {
                "type": "integer",
                "description": "'scroll_down'或'scroll_up'操作的滚动像素数(向下为正,向上为负)",
            },
            "tab_id": {
                "type": "integer",
                "description": "'switch_tab'操作的标签页ID",
            },
            "query": {
                "type": "string",
                "description": "'web_search'操作的搜索查询",
            },
            "goal": {
                "type": "string",
                "description": "'extract_content'操作的提取目标",
            },
            "keys": {
                "type": "string",
                "description": "'send_keys'操作要发送的按键",
            },
            "seconds": {
                "type": "integer",
                "description": "'wait'操作的等待秒数",
            },
        },
        "required": ["action"],
        "dependencies": {
            "go_to_url": ["url"],
            "click_element": ["index"],
            "input_text": ["index", "text"],
            "switch_tab": ["tab_id"],
            "open_tab": ["url"],
            "scroll_down": ["scroll_amount"],
            "scroll_up": ["scroll_amount"],
            "scroll_to_text": ["text"],
            "send_keys": ["keys"],
            "get_dropdown_options": ["index"],
            "select_dropdown_option": ["index", "text"],
            "go_back": [],
            "web_search": ["query"],
            "wait": ["seconds"],
            "extract_content": ["goal"],
        },
    }

    lock: asyncio.Lock = Field(default_factory=asyncio.Lock)
    browser: Optional[BrowserUseBrowser] = Field(default=None, exclude=True)
    context: Optional[BrowserContext] = Field(default=None, exclude=True)
    dom_service: Optional[DomService] = Field(default=None, exclude=True)
    web_search_tool: WebSearch = Field(default_factory=WebSearch, exclude=True)

    # 通用功能的上下文
    tool_context: Optional[Context] = Field(default=None, exclude=True)

    llm: Optional[LLM] = Field(default_factory=LLM)

    @field_validator("parameters", mode="before")
    def validate_parameters(cls, v: dict, info: ValidationInfo) -> dict:
        if not v:
            raise ValueError("参数不能为空")
        return v

    async def _ensure_browser_initialized(self) -> BrowserContext:
        """确保浏览器和上下文已初始化。"""
        if self.browser is None:
            browser_config_kwargs = {"headless": False, "disable_security": True}

            if config.browser_config:
                from browser_use.browser.browser import ProxySettings

                # 处理代理设置
                if config.browser_config.proxy and config.browser_config.proxy.server:
                    browser_config_kwargs["proxy"] = ProxySettings(
                        server=config.browser_config.proxy.server,
                        username=config.browser_config.proxy.username,
                        password=config.browser_config.proxy.password,
                    )

                browser_attrs = [
                    "headless",
                    "disable_security",
                    "extra_chromium_args",
                    "chrome_instance_path",
                    "wss_url",
                    "cdp_url",
                ]

                for attr in browser_attrs:
                    value = getattr(config.browser_config, attr, None)
                    if value is not None:
                        if not isinstance(value, list) or value:
                            browser_config_kwargs[attr] = value

            self.browser = BrowserUseBrowser(BrowserConfig(**browser_config_kwargs))

        if self.context is None:
            context_config = BrowserContextConfig()

            # 如果配置中有上下文配置,则使用它
            if (
                config.browser_config
                and hasattr(config.browser_config, "new_context_config")
                and config.browser_config.new_context_config
            ):
                context_config = config.browser_config.new_context_config

            self.context = await self.browser.new_context(context_config)
            self.dom_service = DomService(await self.context.get_current_page())

        return self.context

    async def execute(
        self,
        action: str,
        url: Optional[str] = None,
        index: Optional[int] = None,
        text: Optional[str] = None,
        scroll_amount: Optional[int] = None,
        tab_id: Optional[int] = None,
        query: Optional[str] = None,
        goal: Optional[str] = None,
        keys: Optional[str] = None,
        seconds: Optional[int] = None,
        **kwargs,
    ) -> ToolResult:
        """
        执行指定的浏览器操作。

        参数:
            action: 要执行的浏览器操作
            url: 导航或新标签页的URL
            index: 点击或输入操作的元素索引
            text: 输入操作或搜索查询的文本
            scroll_amount: 滚动操作的像素数
            tab_id: switch_tab操作的标签页ID
            query: Google搜索的查询
            goal: 内容提取的目标
            keys: 键盘操作要发送的按键
            seconds: 等待的秒数
            **kwargs: 额外参数

        返回:
            包含操作输出或错误的ToolResult
        """
        async with self.lock:
            try:
                context = await self._ensure_browser_initialized()

                # 从配置获取最大内容长度
                max_content_length = getattr(
                    config.browser_config, "max_content_length", 2000
                )

                # 导航操作
                if action == "go_to_url":
                    if not url:
                        return ToolResult(
                            error="'go_to_url'操作需要URL"
                        )
                    page = await context.get_current_page()
                    await page.goto(url)
                    await page.wait_for_load_state()
                    return ToolResult(output=f"已导航到 {url}")

                elif action == "go_back":
                    await context.go_back()
                    return ToolResult(output="已返回上一页")

                elif action == "refresh":
                    await context.refresh_page()
                    return ToolResult(output="已刷新当前页面")

                elif action == "web_search":
                    if not query:
                        return ToolResult(
                            error="'web_search'操作需要查询"
                        )
                    search_results = await self.web_search_tool.execute(query)

                    if search_results:
                        # 导航到第一个搜索结果
                        first_result = search_results[0]
                        if isinstance(first_result, dict) and "url" in first_result:
                            url_to_navigate = first_result["url"]
                        elif isinstance(first_result, str):
                            url_to_navigate = first_result
                        else:
                            return ToolResult(
                                error=f"无效的搜索结果格式: {first_result}"
                            )

                        page = await context.get_current_page()
                        await page.goto(url_to_navigate)
                        await page.wait_for_load_state()

                        return ToolResult(
                            output=f"已搜索'{query}'并导航到第一个结果: {url_to_navigate}\n所有结果:"
                            + "\n".join([str(r) for r in search_results])
                        )
                    else:
                        return ToolResult(
                            error=f"未找到'{query}'的搜索结果"
                        )

                # 元素交互操作
                elif action == "click_element":
                    if index is None:
                        return ToolResult(
                            error="'click_element'操作需要索引"
                        )
                    element = await context.get_dom_element_by_index(index)
                    if not element:
                        return ToolResult(error=f"未找到索引为{index}的元素")
                    download_path = await context._click_element_node(element)
                    output = f"已点击索引为{index}的元素"
                    if download_path:
                        output += f" - 已下载文件到 {download_path}"
                    return ToolResult(output=output)

                elif action == "input_text":
                    if index is None or not text:
                        return ToolResult(
                            error="'input_text'操作需要索引和文本"
                        )
                    element = await context.get_dom_element_by_index(index)
                    if not element:
                        return ToolResult(error=f"未找到索引为{index}的元素")
                    await context._input_text_element_node(element, text)
                    return ToolResult(
                        output=f"已在索引为{index}的元素中输入'{text}'"
                    )

                elif action == "scroll_down" or action == "scroll_up":
                    direction = 1 if action == "scroll_down" else -1
                    amount = (
                        scroll_amount
                        if scroll_amount is not None
                        else context.config.browser_window_size["height"]
                    )
                    await context.execute_javascript(
                        f"window.scrollBy(0, {direction * amount});"
                    )
                    return ToolResult(
                        output=f"已向{'下' if direction > 0 else '上'}滚动{amount}像素"
                    )

                elif action == "scroll_to_text":
                    if not text:
                        return ToolResult(
                            error="'scroll_to_text'操作需要文本"
                        )
                    page = await context.get_current_page()
                    try:
                        locator = page.get_by_text(text, exact=False)
                        await locator.scroll_into_view_if_needed()
                        return ToolResult(output=f"已滚动到文本: '{text}'")
                    except Exception as e:
                        return ToolResult(error=f"滚动到文本失败: {str(e)}")

                elif action == "send_keys":
                    if not keys:
                        return ToolResult(
                            error="'send_keys'操作需要按键"
                        )
                    page = await context.get_current_page()
                    await page.keyboard.press(keys)
                    return ToolResult(output=f"已发送按键: {keys}")

                elif action == "get_dropdown_options":
                    if index is None:
                        return ToolResult(
                            error="'get_dropdown_options'操作需要索引"
                        )
                    element = await context.get_dom_element_by_index(index)
                    if not element:
                        return ToolResult(error=f"未找到索引为{index}的元素")
                    page = await context.get_current_page()
                    options = await page.evaluate(
                        """
                        (xpath) => {
                            const select = document.evaluate(xpath, document, null,
                                XPathResult.FIRST_ORDERED_NODE_TYPE, null).singleNodeValue;
                            if (!select) return null;
                            return Array.from(select.options).map(opt => ({
                                text: opt.text,
                                value: opt.value,
                                index: opt.index
                            }));
                        }
                    """,
                        element.xpath,
                    )
                    return ToolResult(output=f"下拉选项: {options}")

                elif action == "select_dropdown_option":
                    if index is None or not text:
                        return ToolResult(
                            error="'select_dropdown_option'操作需要索引和文本"
                        )
                    element = await context.get_dom_element_by_index(index)
                    if not element:
                        return ToolResult(error=f"未找到索引为{index}的元素")
                    page = await context.get_current_page()
                    await page.select_option(element.xpath, label=text)
                    return ToolResult(
                        output=f"已在索引为{index}的下拉列表中选择选项'{text}'"
                    )

                # 内容提取操作
                elif action == "extract_content":
                    if not goal:
                        return ToolResult(
                            error="'extract_content'操作需要目标"
                        )
                    page = await context.get_current_page()
                    try:
                        # 获取页面内容并转换为markdown以便更好地处理
                        html_content = await page.content()

                        # 在此处导入markdownify以避免全局导入
                        try:
                            import markdownify

                            content = markdownify.markdownify(html_content)
                        except ImportError:
                            # 如果markdownify不可用则回退
                            content = html_content

                        # 为LLM创建提示词
                        prompt_text = """
你的任务是提取页面内容。你将获得一个页面和一个目标,你应该从页面中提取与此目标相关的所有信息。如果目标模糊,则总结页面。以json格式响应。
提取目标: {goal}

页面内容:
{page}
"""
                        # 用目标和内容格式化提示词
                        max_content_length = min(50000, len(content))
                        formatted_prompt = prompt_text.format(
                            goal=goal, page=content[:max_content_length]
                        )

                        # 为LLM创建适当的消息列表
                        from app.schema import Message

                        messages = [Message.user_message(formatted_prompt)]

                        # 为工具定义提取函数
                        extraction_function = {
                            "type": "function",
                            "function": {
                                "name": "extract_content",
                                "description": "根据目标从网页提取特定信息",
                                "parameters": {
                                    "type": "object",
                                    "properties": {
                                        "extracted_content": {
                                            "type": "object",
                                            "description": "根据目标从页面提取的内容",
                                            "properties": {
                                                "text": {
                                                    "type": "string",
                                                    "description": "从页面提取的文本内容",
                                                },
                                                "metadata": {
                                                    "type": "object",
                                                    "description": "关于提取内容的额外元数据",
                                                    "properties": {
                                                        "source": {
                                                            "type": "string",
                                                            "description": "提取内容的来源",
                                                        }
                                                    },
                                                },
                                            },
                                        }
                                    },
                                    "required": ["extracted_content"],
                                },
                            },
                        }

                        # 使用LLM提取内容并要求函数调用
                        response = await self.llm.ask_tool(
                            messages,
                            tools=[extraction_function],
                            tool_choice="required",
                        )

                        # 从函数调用响应中提取内容
                        if (
                            response
                            and response.tool_calls
                            and len(response.tool_calls) > 0
                        ):
                            # 获取第一个工具调用参数
                            tool_call = response.tool_calls[0]
                            # 解析JSON参数
                            try:
                                args = json.loads(tool_call.function.arguments)
                                extracted_content = args.get("extracted_content", {})
                                # 将提取的内容格式化为JSON字符串
                                content_json = json.dumps(
                                    extracted_content, indent=2, ensure_ascii=False
                                )
                                msg = f"从页面提取的内容:\n{content_json}\n"
                            except Exception as e:
                                msg = f"解析提取结果时出错: {str(e)}\n原始响应: {tool_call.function.arguments}"
                        else:
                            msg = "未从页面提取到内容。"

                        return ToolResult(output=msg)
                    except Exception as e:
                        # 提供更有帮助的错误消息
                        error_msg = f"提取内容失败: {str(e)}"
                        try:
                            # 尝试作为回退返回部分页面内容
                            return ToolResult(
                                output=f"{error_msg}\n这是页面内容的一部分:\n{content[:2000]}..."
                            )
                        except:
                            # 如果所有尝试都失败,只返回错误
                            return ToolResult(error=error_msg)

                # 标签页管理操作
                elif action == "switch_tab":
                    if tab_id is None:
                        return ToolResult(
                            error="'switch_tab'操作需要标签页ID"
                        )
                    await context.switch_to_tab(tab_id)
                    page = await context.get_current_page()
                    await page.wait_for_load_state()
                    return ToolResult(output=f"已切换到标签页{tab_id}")

                elif action == "open_tab":
                    if not url:
                        return ToolResult(error="'open_tab'操作需要URL")
                    await context.create_new_tab(url)
                    return ToolResult(output=f"已用{url}打开新标签页")

                elif action == "close_tab":
                    await context.close_current_tab()
                    return ToolResult(output="已关闭当前标签页")

                # 实用工具操作
                elif action == "wait":
                    seconds_to_wait = seconds if seconds is not None else 3
                    await asyncio.sleep(seconds_to_wait)
                    return ToolResult(output=f"已等待{seconds_to_wait}秒")

                else:
                    return ToolResult(error=f"未知操作: {action}")

            except Exception as e:
                return ToolResult(error=f"浏览器操作'{action}'失败: {str(e)}")

    async def get_current_state(
        self, context: Optional[BrowserContext] = None
    ) -> ToolResult:
        """
        获取当前浏览器状态作为ToolResult。
        如果未提供context,则使用self.context。
        """
        try:
            # 使用提供的context或回退到self.context
            ctx = context or self.context
            if not ctx:
                return ToolResult(error="浏览器上下文未初始化")

            state = await ctx.get_state()

            # 如果viewport_info不存在则创建
            viewport_height = 0
            if hasattr(state, "viewport_info") and state.viewport_info:
                viewport_height = state.viewport_info.height
            elif hasattr(ctx, "config") and hasattr(ctx.config, "browser_window_size"):
                viewport_height = ctx.config.browser_window_size.get("height", 0)

            # 为状态截取屏幕截图
            page = await ctx.get_current_page()

            await page.bring_to_front()
            await page.wait_for_load_state()

            screenshot = await page.screenshot(
                full_page=True, animations="disabled", type="jpeg", quality=100
            )

            screenshot = base64.b64encode(screenshot).decode("utf-8")

            # 构建包含所有必需字段的状态信息
            state_info = {
                "url": state.url,
                "title": state.title,
                "tabs": [tab.model_dump() for tab in state.tabs],
                "help": "[0]、[1]、[2]等表示可点击的索引,对应列出的元素。点击这些索引将导航到或与相应内容交互。",
                "interactive_elements": (
                    state.element_tree.clickable_elements_to_string()
                    if state.element_tree
                    else ""
                ),
                "scroll_info": {
                    "pixels_above": getattr(state, "pixels_above", 0),
                    "pixels_below": getattr(state, "pixels_below", 0),
                    "total_height": getattr(state, "pixels_above", 0)
                    + getattr(state, "pixels_below", 0)
                    + viewport_height,
                },
                "viewport_height": viewport_height,
            }

            return ToolResult(
                output=json.dumps(state_info, indent=4, ensure_ascii=False),
                base64_image=screenshot,
            )
        except Exception as e:
            return ToolResult(error=f"获取浏览器状态失败: {str(e)}")

    async def cleanup(self):
        """清理浏览器资源。"""
        async with self.lock:
            if self.context is not None:
                await self.context.close()
                self.context = None
                self.dom_service = None
            if self.browser is not None:
                await self.browser.close()
                self.browser = None

    def __del__(self):
        """确保对象销毁时进行清理。"""
        if self.browser is not None or self.context is not None:
            try:
                asyncio.run(self.cleanup())
            except RuntimeError:
                loop = asyncio.new_event_loop()
                loop.run_until_complete(self.cleanup())
                loop.close()

    @classmethod
    def create_with_context(cls, context: Context) -> "BrowserUseTool[Context]":
        """用特定上下文创建BrowserUseTool的工厂方法。"""
        tool = cls()
        tool.tool_context = context
        return tool
